// 融合显示程序
// 功能：控制松灵小车运动，显示相机图像、激光雷达点云和融合后的图像

#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/PointCloud2.h>
#include <cv_bridge/cv_bridge.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/visualization/cloud_viewer.h>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>
#include <iostream>
#include <vector>

class FusionDisplay {
private:
    ros::NodeHandle nh_;
    ros::Subscriber image_sub_;
    ros::Subscriber cloud_sub_;
    ros::Publisher fusion_image_pub_;
    ros::Publisher rgb_cloud_pub_;
    
    // 数据存储
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_;
    cv::Mat current_image_;
    
    // 标定参数
    cv::Mat camera_matrix_;
    cv::Mat distortion_coeffs_;
    cv::Mat rotation_matrix_;
    cv::Mat translation_vector_;
    
    // 显示窗口
    static pcl::visualization::CloudViewer pcl_viewer_;
    static const std::string origin_image_window_;
    static const std::string fused_image_window_;
    
    // 颜色映射
    int color_map_[21][3];
    float color_distance_;

public:
    FusionDisplay() : 
        cloud_(new pcl::PointCloud<pcl::PointXYZ>()),
        color_distance_(1.2) {
        
        // 初始化颜色映射
        initColorMap();
        
        // 加载标定参数
        loadCalibrationParameters();
        
        // 设置话题订阅和发布
        image_sub_ = nh_.subscribe("/camera/color/image_raw", 10, 
                                 &FusionDisplay::imageCallback, this);
        cloud_sub_ = nh_.subscribe("/rslidar_points", 10,
                                 &FusionDisplay::cloudCallback, this);
        fusion_image_pub_ = nh_.advertise<sensor_msgs::Image>("/fusion_image", 10);
        rgb_cloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("/rgb_cloud", 10);
        
        // 创建显示窗口
        cv::namedWindow(origin_image_window_, cv::WINDOW_AUTOSIZE);
        cv::namedWindow(fused_image_window_, cv::WINDOW_AUTOSIZE);
        
        ROS_INFO("Fusion Display initialized");
        ROS_INFO("Displaying: Original Image, Point Cloud, and Fused Image");
        ROS_INFO("Use remote control to move the robot and observe the fusion results");
    }
    
    void initColorMap() {
        int temp_colors[21][3] = {
            {255, 0, 0}, {255, 69, 0}, {255, 99, 71}, {255, 140, 0}, {255, 165, 0},
            {238, 173, 14}, {255, 193, 37}, {255, 255, 0}, {255, 236, 139}, {202, 255, 112},
            {0, 255, 0}, {84, 255, 159}, {127, 255, 212}, {0, 229, 238}, {152, 245, 255},
            {178, 223, 238}, {126, 192, 238}, {28, 134, 238}, {0, 0, 255}, {72, 118, 255}, {122, 103, 238}
        };
        
        for (int i = 0; i < 21; i++) {
            for (int j = 0; j < 3; j++) {
                color_map_[i][j] = temp_colors[i][j];
            }
        }
    }
    
    void loadCalibrationParameters() {
        // 尝试加载外参标定结果
        cv::FileStorage fs("camera_lidar_extrinsic_calibration.yaml", cv::FileStorage::READ);
        if (fs.isOpened()) {
            cv::Mat extrinsic_matrix;
            fs["CameraExtrinsicMat"] >> extrinsic_matrix;
            fs["CameraMat"] >> camera_matrix_;
            fs["DistCoeff"] >> distortion_coeffs_;
            fs.release();
            
            // 提取旋转矩阵和平移向量
            rotation_matrix_ = extrinsic_matrix(cv::Rect(0, 0, 3, 3));
            translation_vector_ = extrinsic_matrix(cv::Rect(3, 0, 1, 3));
            
            ROS_INFO("Calibration parameters loaded successfully");
        } else {
            ROS_WARN("Calibration file not found, using default parameters");
            // 使用默认参数
            camera_matrix_ = (cv::Mat_<double>(3, 3) << 
                587.0, 0.0, 320.0,
                0.0, 587.0, 240.0,
                0.0, 0.0, 1.0);
            distortion_coeffs_ = (cv::Mat_<double>(1, 5) << 0.0, 0.0, 0.0, 0.0, 0.0);
            rotation_matrix_ = cv::Mat::eye(3, 3, CV_64F);
            translation_vector_ = (cv::Mat_<double>(3, 1) << 0.0, 0.0, 0.0);
        }
    }

    void imageCallback(const sensor_msgs::Image::ConstPtr& msg) {
        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return;
        }
        
        current_image_ = cv_ptr->image.clone();
        
        // 显示原始图像
        cv::imshow(origin_image_window_, current_image_);
        
        // 执行融合
        performFusion();
        
        // 检查退出条件
        if (cv::waitKey(1) == 27) { // ESC键
            ros::shutdown();
        }
    }
    
    void cloudCallback(const sensor_msgs::PointCloud2::ConstPtr& msg) {
        pcl::fromROSMsg(*msg, *cloud_);
        
        // 显示点云
        pcl_viewer_.showCloud(cloud_);
    }
    
    void performFusion() {
        if (cloud_->empty() || current_image_.empty()) {
            return;
        }
        
        // 将点云投影到图像上
        std::vector<cv::Point3f> points3d;
        points3d.reserve(cloud_->size());
        
        for (size_t i = 0; i < cloud_->size(); i++) {
            points3d.push_back(cv::Point3f(cloud_->points[i].x, 
                                          cloud_->points[i].y, 
                                          cloud_->points[i].z));
        }
        
        if (points3d.empty()) {
            ROS_WARN("No valid 3D points for projection");
            return;
        }
        
        // 投影点云到图像平面
        std::vector<cv::Point2f> projected_points;
        cv::Mat rvec;
        cv::Rodrigues(rotation_matrix_, rvec);
        cv::projectPoints(points3d, rvec, translation_vector_, 
                         camera_matrix_, distortion_coeffs_, projected_points);
        
        // 创建融合图像
        cv::Mat fused_image = current_image_.clone();
        
        // 创建带颜色的点云
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr rgb_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
        std::vector<cv::Scalar> colors;
        colors.reserve(cloud_->size());
        
        for (size_t i = 0; i < projected_points.size(); i++) {
            cv::Point2f p = projected_points[i];
            
            // 创建RGB点云点
            pcl::PointXYZRGB rgb_point;
            rgb_point.x = cloud_->points[i].x;
            rgb_point.y = cloud_->points[i].y;
            rgb_point.z = cloud_->points[i].z;
            rgb_point.r = 0;
            rgb_point.g = 0;
            rgb_point.b = 0;
            
            // 检查投影点是否在图像范围内
            if (p.x >= 0 && p.x < current_image_.cols && 
                p.y >= 0 && p.y < current_image_.rows) {
                
                // 从图像中获取颜色
                cv::Vec3b pixel_color = current_image_.at<cv::Vec3b>(p.y, p.x);
                rgb_point.b = pixel_color[0];
                rgb_point.g = pixel_color[1];
                rgb_point.r = pixel_color[2];
                
                // 根据距离设置颜色
                int color_index = std::min(20, (int)(std::abs(rgb_point.z) / (color_distance_ + 0.00001)));
                cv::Scalar distance_color(color_map_[color_index][2], 
                                        color_map_[color_index][1], 
                                        color_map_[color_index][0]);
                colors.push_back(distance_color);
                
                // 在融合图像上绘制点
                cv::circle(fused_image, p, 1, distance_color, -1);
            } else {
                colors.push_back(cv::Scalar(0, 0, 0));
            }
            
            rgb_cloud->push_back(rgb_point);
        }
        
        // 发布RGB点云
        sensor_msgs::PointCloud2 rgb_cloud_msg;
        pcl::toROSMsg(*rgb_cloud, rgb_cloud_msg);
        rgb_cloud_msg.header.frame_id = "rslidar";
        rgb_cloud_pub_.publish(rgb_cloud_msg);
        
        // 发布融合图像
        sensor_msgs::ImagePtr fusion_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", fused_image).toImageMsg();
        fusion_image_pub_.publish(fusion_msg);
        
        // 显示融合图像
        cv::imshow(fused_image_window_, fused_image);
        
        ROS_INFO("Fusion completed: %zu points projected", projected_points.size());
    }
};

// 静态成员定义
pcl::visualization::CloudViewer FusionDisplay::pcl_viewer_("Point Cloud Viewer");
const std::string FusionDisplay::origin_image_window_ = "Original Camera Image";
const std::string FusionDisplay::fused_image_window_ = "Fused Image";

int main(int argc, char** argv) {
    ros::init(argc, argv, "fusion_display");
    
    ROS_INFO("=== Scout Mini Robot Fusion Display ===");
    ROS_INFO("This program displays:");
    ROS_INFO("1. Original camera image");
    ROS_INFO("2. LiDAR point cloud");
    ROS_INFO("3. Fused point-image");
    ROS_INFO("Use the remote control to move the robot and observe the fusion results");
    ROS_INFO("Press ESC to exit");
    
    FusionDisplay display;
    
    ros::spin();
    
    return 0;
}
